{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    DiagramBuilder,\n",
    "    MeshcatVisualizer,\n",
    "    RigidTransform,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import AddShape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic RRT\n",
    "\n",
    "Note that I've inserted a `sleep` command in the visualization to slow things down so you can watch the tree grow.\n",
    "\n",
    "TODO(russt): Consider adding the voronoi visualization, but it would add a dependency on scipy.  (That's a big dependency for a little example!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def basic_rrt():\n",
    "    N = 10000 if running_as_notebook else 3\n",
    "    Q = np.empty((N, 2))\n",
    "    rng = np.random.default_rng()\n",
    "    Q[0] = rng.random((1, 2))\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=0, xmax=1, ymin=0, ymax=1)\n",
    "    meshcat.StartRecording(set_visualizations_while_recording=False)\n",
    "\n",
    "    start = np.empty((N, 3))\n",
    "    end = np.empty((N, 3))\n",
    "    last_plotted_index = 0\n",
    "    plot_count = 1\n",
    "    for n in range(1, N):\n",
    "        q_sample = rng.random((1, 2))[0]\n",
    "        distance_sq = np.sum((Q[:n] - q_sample) ** 2, axis=1)\n",
    "        closest = np.argmin(distance_sq)\n",
    "        distance = np.sqrt(distance_sq[closest])\n",
    "        if distance > 0.1:\n",
    "            q_sample = Q[closest] + (0.1 / distance) * (q_sample - Q[closest])\n",
    "        start[n - 1] = [Q[closest, 0], 0, Q[closest, 1]]\n",
    "        end[n - 1] = [q_sample[0], 0, q_sample[1]]\n",
    "        if (n < 1000 and n % 100 == 1) or n % 1000 == 1:\n",
    "            meshcat.SetLineSegments(\n",
    "                f\"rrt/{last_plotted_index}\",\n",
    "                start[last_plotted_index:n].T,\n",
    "                end[last_plotted_index:n].T,\n",
    "            )\n",
    "            meshcat.SetProperty(\n",
    "                f\"rrt/{last_plotted_index}\", \"visible\", False, 0\n",
    "            )\n",
    "            meshcat.SetProperty(\n",
    "                f\"rrt/{last_plotted_index}\", \"visible\", True, 0.25 * plot_count\n",
    "            )\n",
    "            plot_count = plot_count + 1\n",
    "            last_plotted_index = n\n",
    "        Q[n] = q_sample\n",
    "\n",
    "    meshcat.PublishRecording()\n",
    "\n",
    "\n",
    "basic_rrt()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RRT Bug trap\n",
    "\n",
    "For bonus points, I'll use SceneGraph for the collision checking.\n",
    "\n",
    "TODO(russt):\n",
    "- Take bigger steps, but check collisions at subsamples along an edge.\n",
    "- Add a goal + goal-bias\n",
    "- Make a version where the robot has geometry, and the collision checks call `plant.SetPosition()`, then `query.HasCollisions()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rrt_bugtrap():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    thickness = 0.05\n",
    "    MITred = [0.6, 0.2, 0.2, 1]\n",
    "    wall = AddShape(plant, Box(0.8, 1.0, thickness), \"bottom\", color=MITred)\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"bottom\", wall),\n",
    "        RigidTransform([0.5, 0, 0.1 + thickness / 2]),\n",
    "    )\n",
    "    wall = AddShape(plant, Box(0.8, 1.0, thickness), \"top\", color=MITred)\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"top\", wall),\n",
    "        RigidTransform([0.5, 0, 0.9 - thickness / 2]),\n",
    "    )\n",
    "    wall = AddShape(\n",
    "        plant, Box(thickness, 1.0, 0.8 - thickness), \"left\", color=MITred\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"left\", wall),\n",
    "        RigidTransform([0.1 + thickness / 2, 0, 0.5]),\n",
    "    )\n",
    "    wall = AddShape(\n",
    "        plant, Box(thickness, 1.0, 0.34), \"right_top\", color=MITred\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"right_top\", wall),\n",
    "        RigidTransform([0.9 - thickness / 2, 0, 0.9 - 0.17]),\n",
    "    )\n",
    "    wall = AddShape(\n",
    "        plant, Box(thickness, 1.0, 0.34), \"right_bottom\", color=MITred\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"right_bottom\", wall),\n",
    "        RigidTransform([0.9 - thickness / 2, 0, 0.1 + 0.17]),\n",
    "    )\n",
    "    wall = AddShape(plant, Box(0.36, 1.0, thickness), \"trap_top\", color=MITred)\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"trap_top\", wall),\n",
    "        RigidTransform([0.9 - 0.18, 0, 0.9 - thickness / 2 - 0.33]),\n",
    "    )\n",
    "    wall = AddShape(\n",
    "        plant, Box(0.36, 1.0, thickness), \"trap_bottom\", color=MITred\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"trap_bottom\", wall),\n",
    "        RigidTransform([0.9 - 0.18, 0, 0.1 + thickness / 2 + 0.33]),\n",
    "    )\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=0, xmax=1, ymin=0, ymax=1)\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.StartRecording(set_visualizations_while_recording=False)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    diagram.ForcedPublish(context)\n",
    "    query = scene_graph.get_query_output_port().Eval(\n",
    "        scene_graph.GetMyContextFromRoot(context)\n",
    "    )\n",
    "\n",
    "    q_init = [0.3, 0.3]\n",
    "\n",
    "    N = 10000 if running_as_notebook else 3\n",
    "    Q = np.empty((N, 2))\n",
    "    rng = np.random.default_rng()\n",
    "    Q[0] = q_init\n",
    "\n",
    "    start = np.empty((N, 3))\n",
    "    end = np.empty((N, 3))\n",
    "\n",
    "    max_length = thickness / 4\n",
    "    n = 1\n",
    "    last_plotted_index = 0\n",
    "    plot_count = 1\n",
    "    while n < N:\n",
    "        q_sample = rng.random((1, 2))[0]\n",
    "        distance_sq = np.sum((Q[:n] - q_sample) ** 2, axis=1)\n",
    "        closest = np.argmin(distance_sq)\n",
    "        distance = np.sqrt(distance_sq[closest])\n",
    "        if distance > max_length:\n",
    "            q_sample = Q[closest] + (max_length / distance) * (\n",
    "                q_sample - Q[closest]\n",
    "            )\n",
    "        if query.ComputeSignedDistanceToPoint(\n",
    "            [q_sample[0], 0, q_sample[1]], 0.0\n",
    "        ):\n",
    "            # Then the sample point is in collision...\n",
    "            continue\n",
    "        start[n - 1] = [Q[closest, 0], 0, Q[closest, 1]]\n",
    "        end[n - 1] = [q_sample[0], 0, q_sample[1]]\n",
    "        if (n < 1000 and n % 100 == 1) or n % 1000 == 1:\n",
    "            meshcat.SetLineSegments(\n",
    "                f\"rrt/{last_plotted_index}\",\n",
    "                start[last_plotted_index:n].T,\n",
    "                end[last_plotted_index:n].T,\n",
    "            )\n",
    "            meshcat.SetProperty(\n",
    "                f\"rrt/{last_plotted_index}\", \"visible\", False, 0\n",
    "            )\n",
    "            meshcat.SetProperty(\n",
    "                f\"rrt/{last_plotted_index}\", \"visible\", True, 0.25 * plot_count\n",
    "            )\n",
    "            plot_count = plot_count + 1\n",
    "            last_plotted_index = n\n",
    "        Q[n] = q_sample\n",
    "        n += 1\n",
    "    meshcat.PublishRecording()\n",
    "\n",
    "\n",
    "rrt_bugtrap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
